#include <iostream>
#include <iomanip>
#include "CudaAESOperations.h"
#include <cuda.h>
#include "aesUtils.h"
#include "cudaUtils.h"


__device__ void cudaCopyPitchedPtr(cudaPitchedPtr dst, cudaPitchedPtr src)
{
    int col = threadIdx.x;
    int row = threadIdx.y;
    byte* destCell = ((byte*) dst.ptr + col * dst.pitch) + row;
    byte* srcCell = ((byte*) src.ptr + col * src.pitch) + row;
    
    *destCell = *srcCell;
}

//state[col][row]
__device__ void cudaByteSub(byte* state)
{
    int col = threadIdx.x;//4..6
    int row = threadIdx.y;
    byte* p = state + col*4 + row;
    byte val = DSBOX[*p];
    *p = val;
}

__global__ void KcudaByteSub(cudaPitchedPtr state)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;

    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* destCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* p = data2 + col*4 + row;
    *p = *destCell;
    cudaByteSub(data2);
    *destCell = *p;
}

__device__ void cudaShiftRow(byte* state, byte* outputState, int nb)
{
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    CUAssert(blockDim.y == 4);

    byte dest_col = DSHIFT_NB[nb][row][col];
    byte* destCell = ((byte*) outputState + (dest_col*4)) + row;
    
    byte* srcCell = state + col*4 + row;
    
    *destCell = *srcCell;
}



__global__ void KcudaShiftRow(cudaPitchedPtr state, cudaPitchedPtr outputState,
        int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch) + row;
    byte* destCell = ((byte*) outputState.ptr + col * outputState.pitch) + row;
        
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    
    cudaShiftRow((byte*)data2, (byte*)buffer2, nb);
    
    *destCell = *(buffer2 + col*4 + row);
}

__device__ void cudaMixColumn(byte* state, byte* outputState, int nb)
{    
    int col = threadIdx.x;
    int row = threadIdx.y;
    //int h = blockDim.y;
    
    byte* destCell = ((byte*) outputState + col * 4) +row;
    byte* srcCell = ((byte*) state + col * 4);
    
    //byte* a = 
    byte a0 = *(srcCell);
    byte a1 = *(srcCell+1);
    byte a2 = *(srcCell+2);
    byte a3 = *(srcCell+3);

    byte tmp, tm; 
    tmp = a0 ^ a1 ^ a2 ^ a3;
    
    byte* myCell = srcCell+row;
    byte* nextCell = srcCell+((row+1)&0x3);//&0x3 == %4 == %h//srcCell+(row+1)%h;
    tm = *myCell ^ *nextCell;
    tm = XTIME(tm);
    *destCell = *myCell ^ tm ^ tmp;    
}

__global__ void KcudaMixColumn(cudaPitchedPtr state, cudaPitchedPtr outputState,
        int nb)
{ 
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch) + row;
    byte* destCell = ((byte*) outputState.ptr + col * outputState.pitch) + row;
        
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    
   // cudaMixColumn(state, outputState, nb);
    cudaMixColumn(data2, buffer2, nb);

    *destCell = *(buffer2 + col*4 + row);
    
}

__device__ void cudaAddRoundKey(byte* state, cudaPitchedPtr key, int nb)
{
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* keyCell = ((byte*) key.ptr + col * key.pitch) + row;
    
    byte* stateCell = ((byte*) state + col * 4) + row;
    
    *stateCell ^= *keyCell;
}


__global__ void KcudaAddRoundKey(cudaPitchedPtr state, cudaPitchedPtr key,
        int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch) + row;
        
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    
    cudaAddRoundKey(data2, key, nb);
    *srcCell = *p;
    
}

__device__ void cudaAesRound(byte* state, byte* cState, cudaPitchedPtr eKey, int nb)
{
    
    cudaByteSub(state);
    // __syncthreads();
    cudaShiftRow((byte*)state, (byte*)cState, nb);
    //__syncthreads();
    // now in cState is right State
    cudaMixColumn(cState, state, nb);
    // __syncthreads();
    //now in state is right state
    cudaAddRoundKey(state, eKey, nb);
    //cudaCopyPitchedPtr(state, eKey);
    
   
}

__global__ void KcudaAesRound(cudaPitchedPtr state, cudaPitchedPtr cState, cudaPitchedPtr eKey, int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch) + row;
        
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    
    cudaAesRound(data2, buffer2, eKey, nb);
    *srcCell = *p;
}

__device__ void cudaAesFinalRound(byte* state, byte* outState, cudaPitchedPtr eKey, int nb)
{
    cudaByteSub(state);
    // __syncthreads();
    cudaShiftRow(state, outState, nb);
    //__syncthreads();
    // now in outState is right State
    cudaAddRoundKey(outState, eKey, nb);
    //cudaCopyPitchedPtr(outState, eKey);
}

__device__ void cudaRijndael(byte* state, byte* outState, cudaPitchedPtr eKey, int nb, int nr)
{
    unsigned int wb = BLOCK_SIZE(nb)/4;
    cudaAddRoundKey(state, eKey, nb);
    
    cudaPitchedPtr tkey;
    tkey.ptr = eKey.ptr;
    tkey.pitch = eKey.pitch;
    tkey.xsize = eKey.xsize;
    tkey.ysize = eKey.ysize;
    tkey.ptr = (void**)((byte*)eKey.ptr +(wb*eKey.pitch));
    
    
    for(int i = 1; i < nr; ++i)
    {
        tkey.ptr = (void**)((byte*)eKey.ptr +(i*wb*eKey.pitch));
        cudaAesRound(state, outState, tkey, nb);
    }
    
    tkey.ptr = (void**)((byte*)eKey.ptr +(nr*wb*eKey.pitch));
    cudaAesFinalRound(state, outState, tkey, nb);
    
}

__global__ void KcudaRijndael(cudaPitchedPtr state, cudaPitchedPtr outState, cudaPitchedPtr eKey, int nb, int nr)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    byte* destCell = ((byte*) state.ptr + col * state.pitch) + row;
    byte* srcCell = ((byte*) outState.ptr + col * outState.pitch) + row;
    byte* pd = data2 + col*4 + row;
    byte* pb = buffer2 + col*4 + row;

    *pd = *destCell;

    cudaRijndael(data2, buffer2, eKey, nb, nr);

    *destCell = *pb ;
    *srcCell = *pb;
    
    /*
    byte* data2 = shared;
    byte* buffer2 = shared+ 32;
    byte* ekey = shared+ 64;

    
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    unsigned int wb = BLOCK_SIZE(nb)/4;
    byte* destCell = ((byte*) state.ptr + col * state.pitch) + row;
    byte* srcCell = ((byte*) outState.ptr + col * outState.pitch) + row;
    byte* keyCell = ((byte*) eKey.ptr + col * eKey.pitch) + row;
    
    byte* pd = data2 + col*4 + row;
    byte* pb = buffer2 + col*4 + row;
    byte* pkey = ekey + col*4 + row;
    

    *pd = *destCell;
    
    for(int i =0; i < nr+1; ++i)
    {
        pkey = ekey + (i*wb + col)*4 + row;
        keyCell = ((byte*) eKey.ptr + (i*wb + col) * eKey.pitch) + row;
        *pkey = *keyCell;
            
    }
    for(int i =0; i <1000; i++)
    cudaAESEncryptBlock(data2, buffer2, ekey, nb, nr);

    *destCell = *pb ;
    *srcCell = *pb;
    */
}



void aesEncryptBlock(byte* state, byte* eKey, const NB nb, const NB nk)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    cudaPitchedPtr d_key;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    d_key.xsize = height;
    d_key.ysize = NR[nk][nb]*width*4;// todo check if it is good size
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_key.ptr,
            &d_key.pitch, d_key.xsize, d_key.ysize));
    
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy2D(d_key.ptr, d_key.pitch, eKey, height, height,
            NR[nk][nb]*width*4, cudaMemcpyHostToDevice));

    
    unsigned int sharedSize = 1000;
    dim3 dimBlock(width, height);//(w, h)
    KcudaRijndael<<<1, dimBlock, sharedSize>>>(d_state, d_outputState, d_key, nb, NR[nk][nb]);
    
    
    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
    CUDA_CHECK(cudaFree(d_key.ptr));
}


void aesRound(byte* state, byte* eKey, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    cudaPitchedPtr d_key;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    d_key.xsize = height;
    d_key.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_key.ptr,
            &d_key.pitch, d_key.xsize, d_key.ysize));
    
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy2D(d_key.ptr, d_key.pitch, eKey, height, height,
            width, cudaMemcpyHostToDevice));

    unsigned int shMemSize = 2 * 32;
    dim3 dimBlock(width, height);//(w, h)
    KcudaAesRound<<<1, dimBlock, shMemSize>>>(d_state, d_outputState, d_key, nb);
    
    
    CUDA_CHECK(cudaMemcpy2D(state, height, d_state.ptr,
            d_state.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
    CUDA_CHECK(cudaFree(d_key.ptr));
}

void mixColumn(byte* state, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);//(w, h)
    KcudaMixColumn<<<1, dimBlock>>>(d_state, d_outputState, nb);

    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
}

void byteSub(byte* state, const NB nb)
{
    cudaPitchedPtr h_state;
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    h_state.xsize = height;
    h_state.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &h_state.ptr, &h_state.pitch,
            h_state.xsize, h_state.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(h_state.ptr, h_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    
    
    dim3 dimBlock(width, height);
    //clock_t start = clock();
    //for (int i =0; i < 1000;++i){
        KcudaByteSub<<<1, dimBlock>>>(h_state);
    //}
    //clock_t end = clock();
    //std::cout << "Cuda Byte Sub time" << (end - start)<< std::endl;
    
    
    
    CUDA_CHECK(cudaMemcpy2D(state, height, h_state.ptr, h_state.pitch, height,
            width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(h_state.ptr));
    
}

void shiftRow(byte* state, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    unsigned int shMemSize = 32*2;
    dim3 dimBlock(width, height);//(w, h)
    KcudaShiftRow<<<1, dimBlock, shMemSize>>>(d_state, d_outputState, nb);

    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
}

void addRoundKey(byte* state, byte* key, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_key;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_key.xsize = height;
    d_key.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_key.ptr, &d_key.pitch, d_key.xsize,
            d_key.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    CUDA_CHECK(cudaMemcpy2D(d_key.ptr, d_key.pitch, key, height, height, width,
            cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);//(w, h)
    KcudaAddRoundKey<<<1, dimBlock>>>(d_state, d_key, nb);

    CUDA_CHECK(cudaMemcpy2D(state, height, d_state.ptr, d_state.pitch, height,
            width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_key.ptr));
    
}


void aesEncryptInit()
{
    //cuInit(0);
    CUDA_CHECK(cudaMemcpyToSymbol(DSBOX, SBOX, 256*sizeof(byte), 0, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpyToSymbol(DSHIFT_NB, SHIFT_NB, SHIFT_NB_SIZE, 0, cudaMemcpyHostToDevice));
}


void expandKey(byte* key, NB keysize, byte * W1, NB nb, int round)
{
    int kw = 0;
    int i = 0;
    int bw = BLOCK_SIZE(nb) / 4;

    kw = (keysize == NB_4) ? 4 : (keysize == NB_6) ? 6 : 8;
    uint32_t* W = reinterpret_cast<uint32_t*> (W1);
    uint32_t temp;

    memcpy((byte*) W, key, kw * 4 * sizeof(byte));
    i = kw;

    if (keysize <= NB_6)
    {
        for (temp = 0; i < bw * (NR[keysize][nb] + 1); ++i)
        {
            temp = W[i - 1];
            if (i % kw == 0)
            {
                temp = subRotByte(temp) ^ RCON[i / kw];
                //temp = subByte(rotByte(temp)) ^ Rcon[i / kw];

            }
            W[i] = W[i - kw] ^ temp;
        }
    }
    else
    {
        for (temp = 0; i < bw * (NR[keysize][nb] + 1); ++i)
        {
            temp = W[i - 1];
            if ((i % kw) == 0)
            {
                temp = subByte(rotByte(temp)) ^ RCON[i / kw];
            }
            else if ((i % kw) == 4)
            {
                temp = subByte(temp);
            }
            W[i] = W[i - kw] ^ temp;
        }
    }

}
